from transformers import ViTImageProcessor, ViTForImageClassification
from transformers import ViTModel
from PIL import Image
import requests


device = "cuda:0"
# device = "cpu"

# vit将这个图片识别为日晷，我不理解，所以我觉得这个东西识别效果极差，不能用于
# image = Image.open('/home/zry/experiments/Switch-NeRF/001938.jpg')
image = Image.open('/home/zry/experiments/Switch-NeRF/auto-sam/output/cmask_image_11.png')

processor = ViTImageProcessor.from_pretrained('/home/zry/experiments/Switch-NeRF/vit/vit-base-patch16-224')
# model = ViTForImageClassification.from_pretrained('/home/zry/experiments/Switch-NeRF/vit/vit-base-patch16-224').to(device)
model = ViTModel.from_pretrained('/home/zry/experiments/Switch-NeRF/vit/vit-base-patch16-224').to(device)

# 这里的input可以输入多batch信息
inputs = processor(images=[image]*2, return_tensors="pt").to(device)
outputs = model(**inputs)
seq = outputs[0][:,0,:] # shape = [batch_size, dim] 这样每个图片都被编码到一个768的向量上去了


logits = outputs.logits

# model predicts one of the 1000 ImageNet classes
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])